{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using RetrieveChat with Qdrant for Retrieve Augmented Code Generation and Question Answering\n",
    "\n",
    "[Qdrant](https://qdrant.tech/) is a high-performance vector search engine/database.\n",
    "\n",
    "This notebook demonstrates the usage of Qdrant for RAG, based on [agentchat_RetrieveChat.ipynb](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat.ipynb).\n",
    "\n",
    "\n",
    "RetrieveChat is a conversational system for retrieve augmented code generation and question answering. In this notebook, we demonstrate how to utilize RetrieveChat to generate code and answer questions based on customized documentations that are not present in the LLM's training dataset. RetrieveChat uses the `AssistantAgent` and `RetrieveUserProxyAgent`, which is similar to the usage of `AssistantAgent` and `UserProxyAgent` in other notebooks (e.g., [Automated Task Solving with Code Generation, Execution & Debugging](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_auto_feedback_from_code_execution.ipynb)).\n",
    "\n",
    "We'll demonstrate usage of RetrieveChat with Qdrant for code generation and question answering w/ human feedback.\n",
    "\n",
    "````{=mdx}\n",
    ":::info Requirements\n",
    "Some extra dependencies are needed for this notebook, which can be installed via pip:\n",
    "\n",
    "```bash\n",
    "pip install \"autogen[retrievechat-qdrant]\" \"flaml[automl]\"\n",
    "```\n",
    "\n",
    "For more information, please refer to the [installation guide](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/installing-ag2).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"autogen[retrievechat-qdrant]\" \"flaml[automl]\" -q"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set your API Endpoint\n",
    "\n",
    "The [`config_list_from_json`](https://docs.ag2.ai/latest/docs/api-reference/autogen/config_list_from_json/#autogen.config_list_from_json) function loads a list of configurations from an environment variable or a json file.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qdrant_client import QdrantClient\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "import autogen\n",
    "from autogen import AssistantAgent\n",
    "from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent\n",
    "\n",
    "# Accepted file formats for that can be stored in\n",
    "# a vector database instance\n",
    "from autogen.retrieve_utils import TEXT_FORMATS\n",
    "\n",
    "config_list = autogen.config_list_from_json(\"OAI_CONFIG_LIST\")\n",
    "\n",
    "assert len(config_list) > 0\n",
    "print(\"models to use: \", [config_list[i][\"model\"] for i in range(len(config_list))])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/llm-configuration).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Accepted file formats for `docs_path`:\")\n",
    "print(TEXT_FORMATS)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct agents for RetrieveChat\n",
    "\n",
    "We start by initializing the `AssistantAgent` and `RetrieveUserProxyAgent`. The system message needs to be set to \"You are a helpful assistant.\" for AssistantAgent. The detailed instructions are given in the user message. Later we will use the `RetrieveUserProxyAgent.generate_init_prompt` to combine the instructions and a retrieval augmented generation task for an initial prompt to be sent to the LLM assistant.\n",
    "\n",
    "### You can find the list of all the embedding models supported by Qdrant [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. create an AssistantAgent instance named \"assistant\"\n",
    "assistant = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    system_message=\"You are a helpful assistant.\",\n",
    "    llm_config={\n",
    "        \"timeout\": 600,\n",
    "        \"cache_seed\": 42,\n",
    "        \"config_list\": config_list,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Optionally create embedding function object\n",
    "sentence_transformer_ef = SentenceTransformer(\"all-distilroberta-v1\").encode\n",
    "client = QdrantClient(\":memory:\")\n",
    "\n",
    "# 2. create the RetrieveUserProxyAgent instance named \"ragproxyagent\"\n",
    "# Refer to https://docs.ag2.ai/latest/docs/api-reference/autogen/agentchat/contrib/retrieve_user_proxy_agent/RetrieveUserProxyAgent/#autogen.agentchat.contrib.retrieve_user_proxy_agent.RetrieveUserProxyAgent\n",
    "# and https://docs.ag2.ai/docs/reference/agentchat/contrib/vectordb/qdrant\n",
    "# for more information on the RetrieveUserProxyAgent and QdrantVectorDB\n",
    "ragproxyagent = RetrieveUserProxyAgent(\n",
    "    name=\"ragproxyagent\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    max_consecutive_auto_reply=10,\n",
    "    retrieve_config={\n",
    "        \"task\": \"code\",\n",
    "        \"docs_path\": [\n",
    "            \"https://raw.githubusercontent.com/ag2ai/flaml/main/README.md\",\n",
    "            \"https://raw.githubusercontent.com/ag2ai/FLAML/main/website/docs/Research.md\",\n",
    "        ],  # change this to your own path, such as https://raw.githubusercontent.com/ag2ai/ag2/main/README.md\n",
    "        \"chunk_token_size\": 2000,\n",
    "        \"model\": config_list[0][\"model\"],\n",
    "        \"db_config\": {\"client\": client},\n",
    "        \"vector_db\": \"qdrant\",  # qdrant database\n",
    "        \"get_or_create\": True,  # set to False if you don't want to reuse an existing collection\n",
    "        \"overwrite\": True,  # set to True if you want to overwrite an existing collection\n",
    "        \"embedding_function\": sentence_transformer_ef,  # If left out fastembed \"BAAI/bge-small-en-v1.5\" will be used\n",
    "    },\n",
    "    code_execution_config=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"example-1\"></a>\n",
    "### Example 1\n",
    "\n",
    "[back to top](#)\n",
    "\n",
    "Use RetrieveChat to answer a question and ask for human-in-loop feedbacks.\n",
    "\n",
    "Problem: Is there a function named `tune_automl` in FLAML?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset the assistant. Always reset the assistant before starting a new conversation.\n",
    "assistant.reset()\n",
    "\n",
    "qa_problem = \"Is there a function called tune_automl?\"\n",
    "chat_results = ragproxyagent.initiate_chat(assistant, message=ragproxyagent.message_generator, problem=qa_problem)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"example-2\"></a>\n",
    "### Example 2\n",
    "\n",
    "[back to top](#)\n",
    "\n",
    "Use RetrieveChat to answer a question that is not related to code generation.\n",
    "\n",
    "Problem: Who is the author of FLAML?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset the assistant. Always reset the assistant before starting a new conversation.\n",
    "assistant.reset()\n",
    "\n",
    "qa_problem = \"Who is the author of FLAML?\"\n",
    "chat_results = ragproxyagent.initiate_chat(assistant, message=ragproxyagent.message_generator, problem=qa_problem)"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "This notebook demonstrates the usage of QdrantRetrieveUserProxyAgent for RAG.",
   "tags": [
    "Qdrant",
    "integration",
    "RAG"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
